# 用于dma测速，写进transfer_rates.csv中

import subprocess
import os
import csv

# 创建结果目录
results_dir = '../results_raw'
os.makedirs(results_dir, exist_ok=True)

# 结果文件
output_file = os.path.join(results_dir, 'transfer_rates.csv')

# 定义不同的参数组合
dma_task_nums = [1, 2, 4, 8, 16, 32, 64, 128]  # 根据需要调整这些值
# 补充完整的dma_task_sizes列表，从8开始的2的幂次方
dma_task_sizes = [2**i for i in range(3, 22)]  # 8 到 2^20的所有2的幂次方

def convert_to_us(value, unit):
    """将时间单位转换为微秒"""
    if unit == 's':
        return value * 1_000_000
    elif unit == 'ms':
        return value * 1_000
    elif unit == 'us':
        return value
    else:
        raise ValueError(f"Unknown unit: {unit}")

def calculate_speed(total_size_bytes, time_us):
    """计算传输速度，单位为M/s"""
    if time_us == 0:
        return float('inf')  # 防止除以零
    speed_mbps = (total_size_bytes / time_us) * 1_000_000 / (1024 * 1024)
    return speed_mbps

# 打开文件以写入结果
with open(output_file, 'w', newline='') as csvfile:
    fieldnames = ['dma_task_num', 'dma_task_size', 'P50_us', 'P99_us', 'Speed_Mbps']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

    writer.writeheader()

    for dma_task_num in dma_task_nums:
        for dma_task_size in dma_task_sizes:
            # 提示当前运行的参数组合
            print(f"Running with dma_task_num={dma_task_num}, dma_task_size={dma_task_size}")

            # 构建命令
            command = [
                './../../bin_dpu/test_communication_client',
                f'--dma_task_num={dma_task_num}',
                f'--dma_task_size={dma_task_size}'
            ]

            # 执行命令并捕获输出
            try:
                result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding='utf-8', errors='ignore')
            except Exception as e:
                print(f"Error executing command: {e}")
                continue

            # 解析输出
            stdout = result.stdout
            lines = stdout.splitlines()
            for line in lines:
                if 'dma_task_num' in line:
                    parts = line.split('|')
                    if len(parts) >= 4:
                        # 提取p50和p99值，并统一单位为us
                        p50_value, p50_unit = parts[2].strip().split()
                        p99_value, p99_unit = parts[3].strip().split()

                        p50_us = convert_to_us(float(p50_value), p50_unit)
                        p99_us = convert_to_us(float(p99_value), p99_unit)

                        # 计算实际传输大小
                        total_size_bytes = dma_task_num * dma_task_size

                        # 计算速度
                        speed_mbps = calculate_speed(total_size_bytes, p50_us)

                        writer.writerow({
                            'dma_task_num': dma_task_num,
                            'dma_task_size': dma_task_size,
                            'P50_us': p50_us,
                            'P99_us': p99_us,
                            'Speed_Mbps': speed_mbps
                        })

print(f"Results saved to {output_file}")
